import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report
from joblib import dump

#读取数据文件
df = pd.read_csv('daily.csv')
#计算收益与风险比率
df['max_ratio'] = df['max_close'] / df['the_close']#收益潜力
df['min_ratio'] = df['min_close'] / df['the_close']#风险程度
#自动划分高低阈值
high_retuen_threshold = df['max_ratio'].quantile(0.4)#前40%定义为高收益
high_risk_threshold = df['min_ratio'].quantile(0.4)#前40%定义为高风险

#生成分类标签
conditions = [
    # (df['max_ratio'] >= high_retuen_threshold ) & (df['min_ratio'] <= high_risk_threshold),
    (df['max_ratio'] >= high_retuen_threshold) & (df['min_ratio'] <= high_risk_threshold),
    (df['max_ratio'] >= high_retuen_threshold ) & (df['min_ratio']> high_risk_threshold),
    (df['max_ratio'] < high_retuen_threshold ) & (df['min_ratio']> high_risk_threshold),
    (df['max_ratio'] < high_retuen_threshold ) & (df['min_ratio']<=high_risk_threshold),
]

labels = ['高收益高风险','高收益低风险','低收益高风险','低收益低风险']
df['category'] = np.select(conditions,labels,default = '未知')

#特征选择
features = df[[
    'eps','total_revenue_ps','undist_profit_ps','gross_margin',
    'fcff','fcfe','tangible_asset','bps','grossprofit_margin','npta','roic'
]]

#标签编码
from sklearn.preprocessing import LabelEncoder
le = LabelEncoder()
df['category_encoded'] = le.fit_transform(df['category'])
#数据标准化
scaler = StandardScaler()
X = scaler.fit_transform(features)
y = df['category_encoded']

#划分训练集和测试集
X_train,X_test,y_train,y_test = train_test_split(X,y,test_size=0.3, ramdom_state=42)

#训练KNN模型
knn = KNeighborsClassifier(n_neighbors=4)
knn.fit(X_train,y_train)


#预测与评估
y_pred = knn.predict(X_test)
print(classification_report(y_test,y_pred, target_names=le.classes_))
